#!/usr/bin/env python3
"""
RocketMQ 基础使用示例 - 电商订单处理
"""

import json
import logging
import time
from typing import Dict, Any


# 假设已导入我们的类
from hy.mq_producer import MessageProducer, MessagePriority
from hy.mq_consumer import MessageConsumer, MessageHandler, ConsumeContext, ConsumeResult

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class OrderProcessor:
    """订单处理器示例"""

    def __init__(self):
        # 初始化消息发送者
        self.producer = MessageProducer()
        logger.info("订单处理器初始化完成")

    def create_order(self, order_data: Dict[str, Any]) -> str:
        """
        创建订单并发送相关消息

        Args:
            order_data: 订单数据

        Returns:
            str: 订单ID
        """
        order_id = f"ORDER_{int(time.time() * 1000)}"

        # 增强订单数据
        enhanced_order_data = {
            "order_id": order_id,
            "user_id": order_data["user_id"],
            "products": order_data["products"],
            "total_amount": order_data["total_amount"],
            "created_time": time.time(),
            "status": "CREATED"
        }

        try:
            # 1. 发送库存扣减消息
            inventory_response = self.producer.send_message(
                topic="INVENTORY_TOPIC",
                body=json.dumps({
                    "order_id": order_id,
                    "products": order_data["products"],
                    "action": "RESERVE"
                }),
                tags="INVENTORY_RESERVE",
                keys=order_id,
                properties={"order_id": order_id, "user_id": str(order_data["user_id"])},
                priority=MessagePriority.HIGH
            )

            if inventory_response.status.name != "SUCCESS":
                raise Exception(f"库存扣减消息发送失败: {inventory_response.error}")

            # 2. 发送支付处理消息
            payment_response = self.producer.send_message(
                topic="PAYMENT_TOPIC",
                body=json.dumps({
                    "order_id": order_id,
                    "user_id": order_data["user_id"],
                    "amount": order_data["total_amount"],
                    "action": "CREATE_PAYMENT"
                }),
                tags="PAYMENT_CREATE",
                keys=order_id,
                properties={"order_id": order_id, "amount": str(order_data["total_amount"])},
                priority=MessagePriority.HIGH
            )

            if payment_response.status.name != "SUCCESS":
                raise Exception(f"支付消息发送失败: {payment_response.error}")

            # 3. 发送订单创建通知（延时消息，用于订单超时处理）
            notification_response = self.producer.send_message(
                topic="ORDER_NOTIFICATION_TOPIC",
                body=json.dumps({
                    "order_id": order_id,
                    "action": "ORDER_TIMEOUT_CHECK",
                    "created_time": enhanced_order_data["created_time"]
                }),
                tags="ORDER_TIMEOUT",
                keys=order_id,
                delay_level=16  # 30分钟后触发
            )

            logger.info(f"订单创建成功: {order_id}")
            logger.info(f"库存消息ID: {inventory_response.message_id}")
            logger.info(f"支付消息ID: {payment_response.message_id}")
            logger.info(f"通知消息ID: {notification_response.message_id}")

            return order_id

        except Exception as e:
            logger.error(f"创建订单失败: {e}")
            raise

    def close(self):
        """关闭资源"""
        self.producer.close()


class InventoryHandler(MessageHandler):
    """库存处理器"""

    def __init__(self):
        # 模拟库存数据库
        self.inventory_db = {
            "PRODUCT_001": 100,
            "PRODUCT_002": 50,
            "PRODUCT_003": 200
        }
        self.reserved_inventory = {}  # 预留库存

    def handle_message(self, context: ConsumeContext) -> ConsumeResult:
        """处理库存相关消息"""
        try:
            # 解析消息内容
            message_body = context.message.body.decode('utf-8')
            data = json.loads(message_body)

            order_id = data["order_id"]
            products = data["products"]
            action = data["action"]

            logger.info(f"处理库存消息: 订单={order_id}, 动作={action}")

            if action == "RESERVE":
                return self._reserve_inventory(order_id, products)
            elif action == "CONFIRM":
                return self._confirm_inventory(order_id, products)
            elif action == "RELEASE":
                return self._release_inventory(order_id, products)
            else:
                logger.warning(f"未知的库存操作: {action}")
                return ConsumeResult.SKIP

        except json.JSONDecodeError as e:
            logger.error(f"消息格式错误: {e}")
            return ConsumeResult.FAILED
        except Exception as e:
            logger.error(f"处理库存消息异常: {e}")
            return ConsumeResult.RETRY

    def _reserve_inventory(self, order_id: str, products: List[Dict]) -> ConsumeResult:
        """预留库存"""
        try:
            # 检查库存是否充足
            for product in products:
                product_id = product["product_id"]
                quantity = product["quantity"]

                if product_id not in self.inventory_db:
                    logger.error(f"产品不存在: {product_id}")
                    return ConsumeResult.FAILED

                available = self.inventory_db[product_id]
                if available < quantity:
                    logger.error(f"库存不足: {product_id}, 需要={quantity}, 可用={available}")
                    return ConsumeResult.FAILED

            # 预留库存
            reserved_products = []
            for product in products:
                product_id = product["product_id"]
                quantity = product["quantity"]

                self.inventory_db[product_id] -= quantity
                reserved_products.append({
                    "product_id": product_id,
                    "quantity": quantity
                })

            self.reserved_inventory[order_id] = reserved_products

            logger.info(f"库存预留成功: {order_id}")
            return ConsumeResult.SUCCESS

        except Exception as e:
            logger.error(f"预留库存失败: {e}")
            return ConsumeResult.RETRY

    def _confirm_inventory(self, order_id: str, products: List[Dict]) -> ConsumeResult:
        """确认库存扣减"""
        if order_id in self.reserved_inventory:
            # 从预留库存中移除
            del self.reserved_inventory[order_id]
            logger.info(f"库存扣减确认: {order_id}")
            return ConsumeResult.SUCCESS
        else:
            logger.warning(f"未找到预留库存: {order_id}")
            return ConsumeResult.SKIP

    def _release_inventory(self, order_id: str, products: List[Dict]) -> ConsumeResult:
        """释放预留库存"""
        if order_id in self.reserved_inventory:
            reserved_products = self.reserved_inventory[order_id]

            # 恢复库存
            for product in reserved_products:
                product_id = product["product_id"]
                quantity = product["quantity"]
                self.inventory_db[product_id] += quantity

            # 移除预留记录
            del self.reserved_inventory[order_id]

            logger.info(f"库存释放成功: {order_id}")
            return ConsumeResult.SUCCESS
        else:
            logger.warning(f"未找到要释放的预留库存: {order_id}")
            return ConsumeResult.SKIP


class PaymentHandler(MessageHandler):
    """支付处理器"""

    def __init__(self):
        # 模拟支付记录
        self.payment_records = {}

    def handle_message(self, context: ConsumeContext) -> ConsumeResult:
        """处理支付相关消息"""
        try:
            message_body = context.message.body.decode('utf-8')
            data = json.loads(message_body)

            order_id = data["order_id"]
            user_id = data["user_id"]
            amount = data["amount"]
            action = data["action"]

            logger.info(f"处理支付消息: 订单={order_id}, 动作={action}, 金额={amount}")

            if action == "CREATE_PAYMENT":
                return self._create_payment(order_id, user_id, amount)
            elif action == "CONFIRM_PAYMENT":
                return self._confirm_payment(order_id)
            elif action == "REFUND_PAYMENT":
                return self._refund_payment(order_id)
            else:
                logger.warning(f"未知的支付操作: {action}")
                return ConsumeResult.SKIP

        except json.JSONDecodeError as e:
            logger.error(f"支付消息格式错误: {e}")
            return ConsumeResult.FAILED
        except Exception as e:
            logger.error(f"处理支付消息异常: {e}")
            return ConsumeResult.RETRY

    def _create_payment(self, order_id: str, user_id: int, amount: float) -> ConsumeResult:
        """创建支付"""
        try:
            # 模拟支付创建逻辑
            payment_id = f"PAY_{int(time.time() * 1000)}"

            # 模拟支付成功率（90%成功）
            import random
            if random.random() < 0.9:
                status = "SUCCESS"
                logger.info(f"支付创建成功: {payment_id}, 订单: {order_id}")
            else:
                status = "FAILED"
                logger.warning(f"支付创建失败: {payment_id}, 订单: {order_id}")

            self.payment_records[order_id] = {
                "payment_id": payment_id,
                "user_id": user_id,
                "amount": amount,
                "status": status,
                "created_time": time.time()
            }

            return ConsumeResult.SUCCESS

        except Exception as e:
            logger.error(f"创建支付失败: {e}")
            return ConsumeResult.RETRY

    def _confirm_payment(self, order_id: str) -> ConsumeResult:
        """确认支付"""
        if order_id in self.payment_records:
            payment_record = self.payment_records[order_id]
            payment_record["status"] = "CONFIRMED"
            payment_record["confirmed_time"] = time.time()

            logger.info(f"支付确认成功: {order_id}")
            return ConsumeResult.SUCCESS
        else:
            logger.warning(f"未找到支付记录: {order_id}")
            return ConsumeResult.FAILED

    def _refund_payment(self, order_id: str) -> ConsumeResult:
        """退款处理"""
        if order_id in self.payment_records:
            payment_record = self.payment_records[order_id]
            payment_record["status"] = "REFUNDED"
            payment_record["refunded_time"] = time.time()

            logger.info(f"退款处理成功: {order_id}")
            return ConsumeResult.SUCCESS
        else:
            logger.warning(f"未找到要退款的支付记录: {order_id}")
            return ConsumeResult.FAILED


def main():
    """主函数 - 演示完整的消息发送和消费流程"""

    # 1. 创建订单处理器
    order_processor = OrderProcessor()

    # 2. 创建库存消费者
    inventory_consumer = MessageConsumer(handler=InventoryHandler())
    inventory_consumer.subscribe("INVENTORY_TOPIC", "INVENTORY_*")

    # 3. 创建支付消费者
    payment_consumer = MessageConsumer(handler=PaymentHandler())
    payment_consumer.subscribe("PAYMENT_TOPIC", "PAYMENT_*")

    try:
        # 启动消费者
        logger.info("启动消费者...")
        inventory_consumer.start_consuming()
        payment_consumer.start_consuming()

        # 模拟订单创建
        time.sleep(2)  # 等待消费者启动

        logger.info("开始创建测试订单...")

        for i in range(5):
            order_data = {
                "user_id": 1001 + i,
                "products": [
                    {"product_id": "PRODUCT_001", "quantity": 2, "price": 99.99},
                    {"product_id": "PRODUCT_002", "quantity": 1, "price": 199.99}
                ],
                "total_amount": 399.97
            }

            order_id = order_processor.create_order(order_data)
            logger.info(f"创建订单 {i + 1}: {order_id}")

            time.sleep(1)  # 间隔1秒

        # 等待消息处理完成
        logger.info("等待消息处理完成...")
        time.sleep(10)

        # 打印统计信息
        logger.info("=== 库存消费者统计 ===")
        inventory_stats = inventory_consumer.get_stats()
        for key, value in inventory_stats.items():
            logger.info(f"{key}: {value}")

        logger.info("=== 支付消费者统计 ===")
        payment_stats = payment_consumer.get_stats()
        for key, value in payment_stats.items():
            logger.info(f"{key}: {value}")

        logger.info("=== 生产者统计 ===")
        producer_stats = order_processor.producer.get_stats()
        for key, value in producer_stats.items():
            logger.info(f"{key}: {value}")

    except KeyboardInterrupt:
        logger.info("接收到中断信号，正在关闭...")
    finally:
        # 关闭资源
        logger.info("关闭资源...")
        inventory_consumer.shutdown()
        payment_consumer.shutdown()
        order_processor.close()

        logger.info("程序结束")


if __name__ == "__main__":
    main()
